package main

import (
	"fmt"
	"iter"
	"os"
	"os/exec"
	"path"

	"github.com/gomlx/gomlx/internal/exceptions"
	"github.com/gomlx/gomlx/internal/must"
	"k8s.io/klog/v2"
)

const outputFile = "gen_builderiface.go"

type config struct {
	hasGraph              bool
	numInputs, numOutputs int
}

func iterConfig() iter.Seq2[int, config] {
	return func(yield func(int, config) bool) {
		var count int
		for _, hasGraph := range []bool{false, true} {
			for _, numInputs := range []int{0, 1, 2, 3, 4, -1} {
				if !hasGraph && numInputs == 0 {
					continue
				}
				for _, numOutputs := range []int{0, 1, 2, 3, 4, -1} {
					if !yield(count, config{hasGraph, numInputs, numOutputs}) {
						return
					}
					count++
				}
			}
		}
	}
}

func main() {
	outputPath := path.Join(must.M1(os.Getwd()), outputFile)
	f := must.M1(os.Create(outputPath))

	w := func(format string, args ...any) {
		_, err := fmt.Fprintf(f, format, args...)
		if err != nil {
			exceptions.Panicf("Failed to write to stdout: %v", err)
		}
	}

	w("package model\n\n")
	w("// ***** This file is generated by internal/cmd/builderiface. DO NOT EDIT DIRECTLY. ***** \n\n")
	w("import (\n")
	w("\t\"github.com/gomlx/gomlx/pkg/core/graph\"\n")
	w("\t\"github.com/pkg/errors\"\n")
	w(")\n\n")

	w("// normalizedBuilderFn is the normalized function used internally by the models.Exec object.\n")
	w("//\n// See convertToNormalizedBuilderFn() to convert from a BuilderFnSet to this type.\n")
	w("type normalizedBuilderFn func (g *graph.Graph, inputs []*graph.Node) ([]*graph.Node)\n\n")

	wFuncType := func(c config) {
		w("func (")
		if c.hasGraph {
			w("g *graph.Graph")
			if c.numInputs != 0 {
				w(", ")
			}
		}
		if c.numInputs < 0 {
			w("inputs ...*graph.Node")
		} else if c.numInputs > 0 {
			for i := range c.numInputs {
				if i > 0 {
					w(", ")
				}
				w("input%d", i)
			}
			w(" *graph.Node")
		}
		w(")")
		if c.numOutputs != 0 {
			w(" (")
			if c.numOutputs != 0 {
				if c.numOutputs < 0 {
					w("outputs []*graph.Node")
				} else {
					for i := range c.numOutputs {
						if i > 0 {
							w(", ")
						}
						w("*graph.Node")
					}
				}
			}
			w(")")
		}
	}

	// Write constraint for any of the build types.
	w("// BuilderFnSet is the set of all valid Build methods signatures supported by the models.Exec object.\n")
	w("type BuilderFnSet interface {\n\t")
	for i, c := range iterConfig() {
		if i > 0 {
			w(" |\n\t")
		}
		wFuncType(c)
	}
	w("\n}\n\n")

	// wInvocation writes the invocation of the builder object, given g, inputs and outputs parameters.
	wInvocation := func(indent string, c config) {

		// Checks the number of inputs passed.
		if c.numInputs >= 0 {
			w("%sif len(inputs) != %d {\n", indent, c.numInputs)
			w("%s\tpanic(errors.Errorf(\"wrong number of inputs for model, expected %d, got %%d\", len(inputs)))\n", indent, c.numInputs)
			w("%s}\n", indent)
		}

		// Assign results to outputs.
		w("%s", indent)
		if c.numOutputs == -1 {
			w("return ")
		} else if c.numOutputs > 0 {
			w("outputs := make([]*graph.Node, %d)\n%s", c.numOutputs, indent)
			for i := range c.numOutputs {
				if i > 0 {
					w(", ")
				}
				w("outputs[%d]", i)
			}
			w(" = ")
		}

		// Write invocation.
		w("builder(")
		if c.hasGraph {
			w("g")
			if c.numInputs != 0 {
				w(", ")
			}
		}
		if c.numInputs < 0 {
			w("inputs...")
		} else if c.numInputs > 0 {
			for i := range c.numInputs {
				if i > 0 {
					w(", ")
				}
				w("inputs[%d]", i)
			}
		}
		w(")")

		// Write return.
		if c.numOutputs == 0 {
			// No outputs.
			w("\n%sreturn nil", indent)
		} else if c.numOutputs > 0 {
			w("\n%sreturn outputs", indent)
		}
		w("\n")
	}

	// Write converter of a model to a generic
	w("// convertToNormalizedBuilderFn converts a build closure compatible with BuilderFnSet to a canonical BuilderFn matching the corresponding parameters.\n" +
		"// It also returns the number of inputs and outputs of the underlying Builder.\n")
	w("func convertToNormalizedBuilderFn[B BuilderFnSet] (builder B) (builderFn normalizedBuilderFn, numInputs, numOutputs int, err error) {\n")
	w("\treturn convertToNormalizedBuilderImpl(builder)\n}\n\n")

	w("func convertToNormalizedBuilderImpl(builderAny any) (builderFn normalizedBuilderFn, numInputs, numOutputs int, err error) {\n")
	w("\tswitch builder := builderAny.(type) {\n")
	for _, c := range iterConfig() {
		w("\tcase ")
		wFuncType(c)
		w(":\n")
		w("\t\treturn func (g *graph.Graph, inputs []*graph.Node) ([]*graph.Node) {\n")
		wInvocation("\t\t\t", c)
		w("\t\t}, %d, %d, nil\n\n", c.numInputs, c.numOutputs)
	}
	w("\tdefault:\n")
	w("\t\treturn nil, 0, 0, errors.Errorf(\"model object passed (%%T) doesn't implement any of the valid Build methods signatures supported, see documentation in models.NewExec for details\", builderAny)\n")
	w("\t}\n}\n")

	must.M(f.Close())
	cmd := exec.Command("go", "fmt", outputPath)
	klog.V(1).Infof("\t%s\n", cmd)
	cmd.Stderr = os.Stderr
	must.M(cmd.Run())

	fmt.Printf("✅ internal/cmd/builderiface:  \tsuccessfully generated %s\n", outputPath)
}
